import torch
import transformers
from transformers import LlamaForCausalLM, LlamaTokenizer
import os

import time
# 计时
start_time = time.time()


# 0:Defalt, all;1:no INFO;2:no WARN;3:no ERROR
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'

# 权重路径 或 "meta-llama/Llama-2-7b-chat-hf"(需要全局VPN下载)
#model_dir = r"./Llama-2-7b-hf"
#model_dir = r"/home/gpu/Llama-2-7b-chat-hf"
model_dir = r"/home/hpdesktop/webui/models/Llama-2-7b-hf"
model = LlamaForCausalLM.from_pretrained(model_dir)        # 初始化模型
tokenizer = LlamaTokenizer.from_pretrained(model_dir)      # 初始化tokenizer

# 使用模型
pipeline = transformers.pipeline(
    "text-generation",
    model=model,
    tokenizer=tokenizer,
    torch_dtype=torch.float16,
    device_map="auto",
)

while True:
    user_input = input("请输入字符串：")
    if user_input == "123":
        print("输入为 '123'，退出循环。")
        break
    # 开始计时
    eval_start = time.time()
    print(f"eval_start_at: {eval_start}")
    # 获取结果
    sequences = pipeline(
        #'I have tomatoes, basil and cheese at home. What can I cook for dinner?\n',
        user_input,
        do_sample=True,
        top_k=10,
        num_return_sequences=2,
        eos_token_id=tokenizer.eos_token_id,
        max_length=200,
    )
    eval_end = time.time()
    print(f"eval_end_at: {eval_end}")
    print(f"eval_cost:{eval_end - eval_start} s")
    # 输出结果
    for seq in sequences:
        print(f"{seq ['generated_text']}")
    # 结束计时
    end_time = time.time()
    elapsed_time = end_time - eval_start

    print(f"总执行时间: {elapsed_time} 秒")
